{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "Notebook for testing the Fenchel-Young loss approach to training what we call PertOpt-net. All code based on "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "import perturbations\n",
    "import fenchel_young as fy\n",
    "from torch_Dijkstra import Dijkstra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Create NN using perturbed differentiable optimization\n",
    "class Pert_ShortestPathNet(nn.Module):\n",
    "    '''\n",
    "    This net is equipped to run an m-by-m grid graphs. No A matrix is necessary.\n",
    "    '''\n",
    "    def __init__(self, m, context_size, device='cpu'):\n",
    "        super().__init__()\n",
    "        self.m = m\n",
    "        self.device = device\n",
    "        self.hidden_dim = 2*context_size\n",
    "\n",
    "        ## Standard layers\n",
    "        self.fc_1 = nn.Linear(context_size, self.hidden_dim)\n",
    "        self.fc_12 = nn.Linear(self.hidden_dim, self.hidden_dim)\n",
    "        self.fc_2 = nn.Linear(self.hidden_dim, self.m**2)\n",
    "        self.leaky_relu = nn.LeakyReLU(0.1)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "      ## Put it all together\n",
    "    def forward(self, d):\n",
    "        w = self.leaky_relu(self.fc_12(self.leaky_relu(self.fc_1(d))))\n",
    "        w = self.relu(self.fc_2(w))\n",
    "        return w.view(w.shape[0], self.m, self.m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cpu'\n",
    "dijkstra = Dijkstra(euclidean_weight=False,four_neighbors=True)\n",
    "criterion2 = nn.MSELoss()\n",
    "criterion = fy.FenchelYoungLoss(dijkstra, num_samples=10, sigma=0.01, noise='gumbel', batched=True, maximize=False, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load data\n",
    "grid_size = 5\n",
    "data_path = '../shortest_path_data/Shortest_Path_training_data'+str(grid_size)+'.pth'\n",
    "state = torch.load(data_path)\n",
    "\n",
    "## Extract data from state\n",
    "train_dataset = state['train_dataset_v']\n",
    "test_dataset = state['test_dataset_v']\n",
    "m = state[\"m\"]\n",
    "A = state[\"A\"].float()\n",
    "b = state[\"b\"].float()\n",
    "num_edges = state[\"num_edges\"]\n",
    "Edge_list = state[\"Edge_list\"]\n",
    "Edge_list_torch = torch.tensor(Edge_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Pert_ShortestPathNet(grid_size, context_size=5, device='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Training setup\n",
    "from torch.utils.data import Dataset, TensorDataset, DataLoader\n",
    "learning_rate = 1e-4\n",
    "test_size = 200\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=200,\n",
    "                              shuffle=True)\n",
    "test_loader = DataLoader(dataset=test_dataset, batch_size=test_size,\n",
    "                             shuffle=False)\n",
    "\n",
    "optimizer = optim.Adam(net.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.1263, 0.0000, 0.0658, 0.1584, 0.0000],\n",
      "        [0.0000, 0.1869, 0.0000, 0.2668, 0.0000],\n",
      "        [0.0000, 0.1536, 0.1520, 0.2397, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2100, 0.2366],\n",
      "        [0.2252, 0.0000, 0.2065, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.43880000710487366\n",
      "tensor([[0.1488, 0.0000, 0.0199, 0.1837, 0.0000],\n",
      "        [0.0000, 0.2182, 0.0000, 0.2441, 0.0000],\n",
      "        [0.0000, 0.2406, 0.1279, 0.2232, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1999, 0.2438],\n",
      "        [0.1761, 0.0000, 0.2053, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.4424000084400177\n",
      "tensor([[0.1343, 0.0000, 0.0433, 0.1853, 0.0000],\n",
      "        [0.0000, 0.2191, 0.0000, 0.2385, 0.0000],\n",
      "        [0.0000, 0.1908, 0.1212, 0.2416, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1888, 0.2394],\n",
      "        [0.1810, 0.0000, 0.1918, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.4580000042915344\n",
      "tensor([[0.1390, 0.0000, 0.0603, 0.2031, 0.0000],\n",
      "        [0.0000, 0.2414, 0.0000, 0.2014, 0.0000],\n",
      "        [0.0000, 0.1631, 0.0952, 0.2511, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1815, 0.2357],\n",
      "        [0.1665, 0.0000, 0.1592, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.438400000333786\n",
      "tensor([[0.1460, 0.0000, 0.0172, 0.1817, 0.0000],\n",
      "        [0.0000, 0.2157, 0.0000, 0.2542, 0.0000],\n",
      "        [0.0000, 0.2461, 0.1280, 0.2229, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1999, 0.2440],\n",
      "        [0.1788, 0.0000, 0.2141, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.4519999921321869\n",
      "tensor([[ 0.0485,  0.0369,  0.0304,  0.0328,  0.0321],\n",
      "        [ 0.0454,  0.0849,  0.0769,  0.0752,  0.0753],\n",
      "        [-0.0047, -0.0035, -0.0085, -0.0007, -0.0009],\n",
      "        [-0.0027, -0.0050, -0.0045, -0.0024, -0.0032],\n",
      "        [-0.0061, -0.0079, -0.0064, -0.0054, -0.0049],\n",
      "        [ 0.0094,  0.0118,  0.0097,  0.0105,  0.0093],\n",
      "        [ 0.0044,  0.0062,  0.0054,  0.0071,  0.0063],\n",
      "        [ 0.0009, -0.0019, -0.0007, -0.0011, -0.0008],\n",
      "        [-0.0290, -0.0084, -0.0309, -0.0131, -0.0222],\n",
      "        [-0.0293, -0.0140, -0.0319, -0.0043, -0.0101]])\n",
      "tensor([ 0.0634,  0.1426, -0.0073, -0.0066, -0.0114,  0.0202,  0.0115, -0.0014,\n",
      "        -0.0408, -0.0335])\n",
      "tensor([[-0.0557, -0.0502, -0.0052,  0.0120,  0.0072,  0.0123,  0.0250,  0.0131,\n",
      "         -0.0814, -0.1056],\n",
      "        [-0.0612, -0.0330, -0.0025,  0.0121,  0.0066,  0.0115,  0.0207,  0.0095,\n",
      "         -0.0433, -0.0760],\n",
      "        [ 0.0174,  0.0056,  0.0014, -0.0026, -0.0016, -0.0028, -0.0038, -0.0018,\n",
      "          0.0004,  0.0140],\n",
      "        [ 0.0680,  0.0322, -0.0015, -0.0154, -0.0085, -0.0144, -0.0311, -0.0134,\n",
      "          0.0677,  0.1214],\n",
      "        [ 0.0378,  0.0210,  0.0004, -0.0068, -0.0044, -0.0081, -0.0168, -0.0078,\n",
      "          0.0405,  0.0754],\n",
      "        [ 0.0072,  0.0106,  0.0010, -0.0022, -0.0012, -0.0025, -0.0063, -0.0029,\n",
      "          0.0228,  0.0278],\n",
      "        [ 0.0340,  0.0286,  0.0017, -0.0062, -0.0034, -0.0069, -0.0120, -0.0060,\n",
      "          0.0397,  0.0493],\n",
      "        [-0.0097, -0.0077, -0.0008,  0.0020,  0.0012,  0.0021,  0.0043,  0.0021,\n",
      "         -0.0126, -0.0181],\n",
      "        [ 0.0068,  0.0139,  0.0012, -0.0029, -0.0016, -0.0026, -0.0078, -0.0039,\n",
      "          0.0328,  0.0347],\n",
      "        [-0.0053, -0.0030, -0.0001,  0.0011,  0.0006,  0.0011,  0.0021,  0.0009,\n",
      "         -0.0047, -0.0078]])\n",
      "tensor([-0.2724, -0.2208,  0.0359,  0.3248,  0.1607,  0.0724,  0.1269, -0.0464,\n",
      "         0.0905, -0.0225])\n",
      "tensor([[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.1203,  0.1459,  0.1015,  0.1259,  0.0145,  0.0096,  0.1060, -0.0112,\n",
      "          0.0773, -0.0158],\n",
      "        [ 0.0961,  0.1318,  0.0954,  0.1074,  0.0142,  0.0095,  0.0856, -0.0096,\n",
      "          0.0736, -0.0138],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0705,  0.0681,  0.0536,  0.0829,  0.0086,  0.0015,  0.0826, -0.0048,\n",
      "          0.0265, -0.0094],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0951,  0.1300,  0.0943,  0.1061,  0.0142,  0.0087,  0.0854, -0.0095,\n",
      "          0.0725, -0.0137],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0408,  0.0470,  0.0355,  0.0538,  0.0079,  0.0012,  0.0475, -0.0026,\n",
      "          0.0180, -0.0052],\n",
      "        [ 0.0676,  0.0413,  0.0245,  0.0572,  0.0028, -0.0007,  0.0669, -0.0046,\n",
      "          0.0120, -0.0068],\n",
      "        [ 0.0352,  0.0289,  0.0172,  0.0352,  0.0033,  0.0019,  0.0323, -0.0024,\n",
      "          0.0085, -0.0034],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [-0.0295, -0.0172, -0.0107, -0.0295, -0.0042,  0.0021, -0.0346,  0.0015,\n",
      "         -0.0016,  0.0028],\n",
      "        [ 0.0635,  0.1047,  0.0803,  0.0749,  0.0109,  0.0075,  0.0572, -0.0074,\n",
      "          0.0659, -0.0109],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [-0.0318, -0.0866, -0.0687, -0.0433, -0.0066, -0.0097, -0.0195,  0.0058,\n",
      "         -0.0643,  0.0078],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000]])\n",
      "tensor([ 0.0000,  0.0000,  0.6400,  0.5500,  0.0000,  0.0000,  0.3600,  0.0000,\n",
      "         0.5450,  0.0000,  0.0000,  0.2125,  0.2800,  0.1500,  0.0000,  0.0000,\n",
      "         0.0000,  0.0000, -0.1195,  0.4150,  0.0000,  0.0000, -0.2855,  0.0000,\n",
      "         0.0000])\n",
      "epoch:  0 | ave_tr_loss:  1.01e-01 | te_loss:  3.16e-01 | acc.:  0.000000 | lr:  1.00e-04 | time:  25.813001      \n",
      "tensor([[0.1688, 0.0000, 0.0574, 0.1947, 0.0000],\n",
      "        [0.0000, 0.2485, 0.0000, 0.1427, 0.0000],\n",
      "        [0.0000, 0.1306, 0.1102, 0.2398, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1709, 0.2363],\n",
      "        [0.1421, 0.0000, 0.1105, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.4440000057220459\n",
      "tensor([[0.1665, 0.0000, 0.0805, 0.2512, 0.0000],\n",
      "        [0.0000, 0.3071, 0.0000, 0.1054, 0.0000],\n",
      "        [0.0000, 0.0786, 0.0588, 0.2768, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1135, 0.2416],\n",
      "        [0.0933, 0.0000, 0.0663, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.44359999895095825\n",
      "tensor([[0.1441, 0.0000, 0.0671, 0.1988, 0.0000],\n",
      "        [0.0000, 0.2429, 0.0000, 0.1814, 0.0000],\n",
      "        [0.0000, 0.1532, 0.0913, 0.2473, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1912, 0.2276],\n",
      "        [0.1691, 0.0000, 0.1472, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.4528000056743622\n",
      "tensor([[0.1555, 0.0000, 0.0285, 0.1736, 0.0000],\n",
      "        [0.0000, 0.2194, 0.0000, 0.2136, 0.0000],\n",
      "        [0.0000, 0.2211, 0.1237, 0.2188, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2102, 0.2321],\n",
      "        [0.1772, 0.0000, 0.1846, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.4359999895095825\n",
      "tensor([[0.1399, 0.0000, 0.1039, 0.2273, 0.0000],\n",
      "        [0.0000, 0.2539, 0.0000, 0.1599, 0.0000],\n",
      "        [0.0000, 0.0407, 0.0821, 0.2840, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1342, 0.2485],\n",
      "        [0.1500, 0.0000, 0.1047, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.4503999948501587\n",
      "tensor([[ 0.0444,  0.0313,  0.0267,  0.0294,  0.0262],\n",
      "        [ 0.0476,  0.0832,  0.0724,  0.0654,  0.0718],\n",
      "        [-0.0052, -0.0052, -0.0096, -0.0011, -0.0016],\n",
      "        [-0.0023, -0.0058, -0.0051, -0.0020, -0.0045],\n",
      "        [-0.0056, -0.0073, -0.0061, -0.0052, -0.0051],\n",
      "        [ 0.0093,  0.0112,  0.0095,  0.0101,  0.0095],\n",
      "        [ 0.0050,  0.0069,  0.0055,  0.0075,  0.0066],\n",
      "        [ 0.0007, -0.0022, -0.0007, -0.0009, -0.0013],\n",
      "        [-0.0269, -0.0102, -0.0300, -0.0142, -0.0169],\n",
      "        [-0.0241, -0.0085, -0.0302, -0.0041, -0.0052]])\n",
      "tensor([ 0.0544,  0.1326, -0.0099, -0.0076, -0.0109,  0.0191,  0.0120, -0.0015,\n",
      "        -0.0421, -0.0304])\n",
      "tensor([[-5.6494e-02, -4.9046e-02, -4.3465e-03,  1.0780e-02,  6.7298e-03,\n",
      "          1.2715e-02,  2.5064e-02,  1.2718e-02, -7.9651e-02, -1.0999e-01],\n",
      "        [-5.8313e-02, -3.0754e-02, -1.5173e-03,  1.2561e-02,  6.6588e-03,\n",
      "          1.1094e-02,  2.0633e-02,  9.3766e-03, -4.3222e-02, -7.2518e-02],\n",
      "        [ 1.2207e-02,  2.6042e-03,  2.9732e-04, -2.1941e-03, -1.1826e-03,\n",
      "         -1.7421e-03, -2.1536e-03, -9.7401e-04, -2.0287e-03,  4.5656e-03],\n",
      "        [ 7.1719e-02,  3.0276e-02, -2.2847e-03, -1.5120e-02, -8.5046e-03,\n",
      "         -1.5202e-02, -3.2065e-02, -1.3351e-02,  6.6048e-02,  1.2689e-01],\n",
      "        [ 3.9012e-02,  1.8772e-02, -9.7786e-04, -5.7268e-03, -3.9074e-03,\n",
      "         -8.4143e-03, -1.6585e-02, -7.2255e-03,  3.8277e-02,  7.7778e-02],\n",
      "        [ 1.1093e-02,  1.1834e-02,  4.8423e-04, -2.3347e-03, -1.4398e-03,\n",
      "         -3.3481e-03, -7.5500e-03, -3.2843e-03,  2.3470e-02,  3.4395e-02],\n",
      "        [ 3.1492e-02,  2.9297e-02,  2.9767e-03, -5.4896e-03, -3.2281e-03,\n",
      "         -6.5081e-03, -1.1094e-02, -5.8491e-03,  3.8580e-02,  4.7014e-02],\n",
      "        [-9.4660e-03, -7.2787e-03, -5.4887e-04,  1.8215e-03,  1.1093e-03,\n",
      "          2.1232e-03,  4.2035e-03,  2.0281e-03, -1.2098e-02, -1.8127e-02],\n",
      "        [ 1.3035e-02,  1.4488e-02,  6.6569e-04, -3.3053e-03, -2.0143e-03,\n",
      "         -3.8252e-03, -9.5453e-03, -4.4751e-03,  3.2974e-02,  4.2945e-02],\n",
      "        [-5.2718e-03, -2.7797e-03, -3.9817e-05,  1.1967e-03,  6.2220e-04,\n",
      "          1.0710e-03,  2.1418e-03,  9.1596e-04, -4.5973e-03, -7.7180e-03]])\n",
      "tensor([-0.2651, -0.2240,  0.0223,  0.3249,  0.1499,  0.0818,  0.1199, -0.0442,\n",
      "         0.1017, -0.0232])\n",
      "tensor([[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 1.1405e-01,  1.3990e-01,  9.4635e-02,  1.1718e-01,  1.5098e-02,\n",
      "          8.5636e-03,  9.6908e-02, -1.0621e-02,  7.2904e-02, -1.4512e-02],\n",
      "        [ 9.4709e-02,  1.2866e-01,  8.9221e-02,  1.0237e-01,  1.5050e-02,\n",
      "          8.5128e-03,  7.9767e-02, -9.2905e-03,  6.9679e-02, -1.2826e-02],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 7.7908e-02,  6.9796e-02,  5.5159e-02,  8.5992e-02,  6.8897e-03,\n",
      "          6.0827e-04,  9.1462e-02, -5.3571e-03,  2.6680e-02, -1.0456e-02],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 9.4709e-02,  1.2866e-01,  8.9221e-02,  1.0237e-01,  1.5050e-02,\n",
      "          8.5128e-03,  7.9767e-02, -9.2905e-03,  6.9679e-02, -1.2826e-02],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 3.3488e-02,  4.3011e-02,  3.5938e-02,  4.8088e-02,  5.2297e-03,\n",
      "          1.0812e-03,  4.3511e-02, -2.2099e-03,  1.8486e-02, -5.1638e-03],\n",
      "        [ 7.4471e-02,  4.4865e-02,  2.7830e-02,  6.1949e-02,  2.8452e-03,\n",
      "         -9.5076e-04,  7.5109e-02, -5.1421e-03,  1.2676e-02, -7.8560e-03],\n",
      "        [ 2.7644e-02,  2.1996e-02,  1.2280e-02,  2.6526e-02,  3.7926e-03,\n",
      "          4.1685e-04,  2.4998e-02, -1.8448e-03,  5.8002e-03, -2.5613e-03],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [-3.0916e-02, -1.7961e-02, -1.3417e-02, -3.1790e-02, -1.9681e-03,\n",
      "          1.0484e-03, -3.8496e-02,  1.5971e-03, -2.2636e-03,  3.4261e-03],\n",
      "        [ 6.9546e-02,  1.0869e-01,  7.7994e-02,  7.8330e-02,  1.2029e-02,\n",
      "          8.0597e-03,  5.7080e-02, -7.6025e-03,  6.4178e-02, -1.0474e-02],\n",
      "        [ 1.1664e-03,  4.3434e-04,  3.9412e-04,  1.0612e-03, -3.0810e-05,\n",
      "         -3.4153e-05,  1.4825e-03, -5.8912e-05, -5.2318e-07, -1.4130e-04],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [-3.6148e-02, -8.8705e-02, -6.3525e-02, -4.4057e-02, -9.2892e-03,\n",
      "         -9.1444e-03, -1.6273e-02,  5.8485e-03, -6.1616e-02,  6.8384e-03],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00]])\n",
      "tensor([ 0.0000,  0.0000,  0.6000,  0.5250,  0.0000,  0.0000,  0.3950,  0.0000,\n",
      "         0.5250,  0.0000,  0.0000,  0.1960,  0.3150,  0.1150,  0.0000,  0.0000,\n",
      "         0.0000,  0.0000, -0.1340,  0.4200,  0.0050,  0.0000, -0.2760,  0.0000,\n",
      "         0.0000])\n",
      "epoch:  0 | ave_tr_loss:  1.79e-01 | te_loss:  3.16e-01 | acc.:  0.000000 | lr:  1.00e-04 | time:  51.998049      \n",
      "tensor([[0.1255, 0.0000, 0.0687, 0.1625, 0.0000],\n",
      "        [0.0000, 0.1931, 0.0000, 0.2535, 0.0000],\n",
      "        [0.0000, 0.1426, 0.1391, 0.2427, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2066, 0.2328],\n",
      "        [0.2180, 0.0000, 0.1987, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.4472000002861023\n",
      "tensor([[0.1667, 0.0000, 0.0501, 0.1910, 0.0000],\n",
      "        [0.0000, 0.2493, 0.0000, 0.1458, 0.0000],\n",
      "        [0.0000, 0.1938, 0.0894, 0.2230, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2157, 0.2157],\n",
      "        [0.1602, 0.0000, 0.1360, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.4503999948501587\n",
      "tensor([[0.1528, 0.0000, 0.0354, 0.1795, 0.0000],\n",
      "        [0.0000, 0.2306, 0.0000, 0.2003, 0.0000],\n",
      "        [0.0000, 0.2163, 0.1073, 0.2230, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2114, 0.2242],\n",
      "        [0.1736, 0.0000, 0.1764, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.43160000443458557\n",
      "tensor([[0.1530, 0.0000, 0.0798, 0.2337, 0.0000],\n",
      "        [0.0000, 0.2830, 0.0000, 0.1233, 0.0000],\n",
      "        [0.0000, 0.1092, 0.0545, 0.2674, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1501, 0.2267],\n",
      "        [0.1254, 0.0000, 0.0990, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.44760000705718994\n",
      "tensor([[0.1465, 0.0000, 0.0689, 0.1785, 0.0000],\n",
      "        [0.0000, 0.2245, 0.0000, 0.1769, 0.0000],\n",
      "        [0.0000, 0.1423, 0.1083, 0.2379, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2068, 0.2198],\n",
      "        [0.1842, 0.0000, 0.1496, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.44999998807907104\n",
      "tensor([[ 0.0418,  0.0254,  0.0206,  0.0260,  0.0252],\n",
      "        [ 0.0501,  0.0905,  0.0768,  0.0673,  0.0701],\n",
      "        [-0.0051, -0.0049, -0.0114, -0.0009, -0.0024],\n",
      "        [-0.0022, -0.0051, -0.0046, -0.0017, -0.0033],\n",
      "        [-0.0053, -0.0071, -0.0059, -0.0047, -0.0049],\n",
      "        [ 0.0091,  0.0113,  0.0094,  0.0096,  0.0089],\n",
      "        [ 0.0047,  0.0066,  0.0048,  0.0071,  0.0062],\n",
      "        [ 0.0010, -0.0019, -0.0005, -0.0006, -0.0009],\n",
      "        [-0.0291, -0.0170, -0.0347, -0.0163, -0.0194],\n",
      "        [-0.0279, -0.0133, -0.0369, -0.0065, -0.0082]])\n",
      "tensor([ 0.0474,  0.1431, -0.0111, -0.0067, -0.0106,  0.0188,  0.0114, -0.0011,\n",
      "        -0.0470, -0.0371])\n",
      "tensor([[-0.0547, -0.0547, -0.0097,  0.0111,  0.0070,  0.0127,  0.0249,  0.0133,\n",
      "         -0.0844, -0.1084],\n",
      "        [-0.0570, -0.0344, -0.0048,  0.0120,  0.0065,  0.0109,  0.0198,  0.0093,\n",
      "         -0.0448, -0.0712],\n",
      "        [ 0.0155,  0.0049,  0.0009, -0.0027, -0.0015, -0.0023, -0.0029, -0.0014,\n",
      "          0.0002,  0.0077],\n",
      "        [ 0.0687,  0.0300,  0.0003, -0.0148, -0.0084, -0.0149, -0.0315, -0.0132,\n",
      "          0.0642,  0.1248],\n",
      "        [ 0.0290,  0.0140,  0.0001, -0.0044, -0.0030, -0.0068, -0.0137, -0.0059,\n",
      "          0.0315,  0.0651],\n",
      "        [ 0.0120,  0.0152,  0.0032, -0.0026, -0.0017, -0.0036, -0.0078, -0.0037,\n",
      "          0.0263,  0.0358],\n",
      "        [ 0.0346,  0.0348,  0.0054, -0.0064, -0.0038, -0.0073, -0.0123, -0.0067,\n",
      "          0.0438,  0.0507],\n",
      "        [-0.0089, -0.0080, -0.0014,  0.0018,  0.0011,  0.0021,  0.0041,  0.0021,\n",
      "         -0.0127, -0.0177],\n",
      "        [ 0.0095,  0.0147,  0.0027, -0.0029, -0.0018, -0.0033, -0.0086, -0.0042,\n",
      "          0.0321,  0.0389],\n",
      "        [-0.0051, -0.0031, -0.0004,  0.0011,  0.0006,  0.0011,  0.0021,  0.0009,\n",
      "         -0.0047, -0.0076]])\n",
      "tensor([-0.2738, -0.2144,  0.0283,  0.3219,  0.1272,  0.0868,  0.1383, -0.0446,\n",
      "         0.0974, -0.0224])\n",
      "tensor([[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.1136,  0.1440,  0.0998,  0.1183,  0.0112,  0.0107,  0.0957, -0.0106,\n",
      "          0.0760, -0.0151],\n",
      "        [ 0.0918,  0.1310,  0.0935,  0.1019,  0.0113,  0.0107,  0.0767, -0.0091,\n",
      "          0.0721, -0.0132],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0757,  0.0754,  0.0597,  0.0895,  0.0077,  0.0033,  0.0880, -0.0050,\n",
      "          0.0286, -0.0103],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0918,  0.1310,  0.0935,  0.1019,  0.0113,  0.0107,  0.0767, -0.0091,\n",
      "          0.0721, -0.0132],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0323,  0.0471,  0.0402,  0.0504,  0.0057,  0.0020,  0.0431, -0.0021,\n",
      "          0.0208, -0.0054],\n",
      "        [ 0.0714,  0.0401,  0.0217,  0.0568,  0.0027, -0.0008,  0.0687, -0.0047,\n",
      "          0.0090, -0.0068],\n",
      "        [ 0.0237,  0.0181,  0.0093,  0.0217,  0.0024,  0.0003,  0.0205, -0.0016,\n",
      "          0.0042, -0.0021],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [-0.0396, -0.0268, -0.0189, -0.0419, -0.0037,  0.0006, -0.0472,  0.0021,\n",
      "         -0.0049,  0.0043],\n",
      "        [ 0.0681,  0.1130,  0.0842,  0.0802,  0.0088,  0.0104,  0.0563, -0.0075,\n",
      "          0.0679, -0.0111],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [-0.0285, -0.0862, -0.0653, -0.0383, -0.0051, -0.0109, -0.0091,  0.0054,\n",
      "         -0.0630,  0.0069],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000]])\n",
      "tensor([ 0.0000,  0.0000,  0.6100,  0.5250,  0.0000,  0.0000,  0.3900,  0.0000,\n",
      "         0.5250,  0.0000,  0.0000,  0.2010,  0.2850,  0.0950,  0.0000,  0.0000,\n",
      "         0.0000,  0.0000, -0.1730,  0.4300,  0.0000,  0.0000, -0.2570,  0.0000,\n",
      "         0.0000])\n",
      "epoch:  0 | ave_tr_loss:  2.39e-01 | te_loss:  3.16e-01 | acc.:  0.000000 | lr:  1.00e-04 | time:  78.233968      \n",
      "tensor([[0.1421, 0.0000, 0.0263, 0.1738, 0.0000],\n",
      "        [0.0000, 0.2214, 0.0000, 0.2339, 0.0000],\n",
      "        [0.0000, 0.2294, 0.1159, 0.2243, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2078, 0.2285],\n",
      "        [0.1795, 0.0000, 0.2010, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.44440001249313354\n",
      "tensor([[0.1262, 0.0000, 0.0738, 0.2168, 0.0000],\n",
      "        [0.0000, 0.2468, 0.0000, 0.1998, 0.0000],\n",
      "        [0.0000, 0.1101, 0.0859, 0.2741, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1456, 0.2368],\n",
      "        [0.1551, 0.0000, 0.1546, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.44920000433921814\n",
      "tensor([[0.1617, 0.0000, 0.0717, 0.2059, 0.0000],\n",
      "        [0.0000, 0.2599, 0.0000, 0.1308, 0.0000],\n",
      "        [0.0000, 0.0822, 0.0974, 0.2579, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1425, 0.2336],\n",
      "        [0.1298, 0.0000, 0.0973, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.446399986743927\n",
      "tensor([[0.1504, 0.0000, 0.1065, 0.2126, 0.0000],\n",
      "        [0.0000, 0.2594, 0.0000, 0.1265, 0.0000],\n",
      "        [0.0000, 0.0439, 0.0753, 0.2696, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1599, 0.2287],\n",
      "        [0.1530, 0.0000, 0.0881, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.4415999948978424\n",
      "tensor([[0.1429, 0.0000, 0.0758, 0.2062, 0.0000],\n",
      "        [0.0000, 0.2395, 0.0000, 0.1659, 0.0000],\n",
      "        [0.0000, 0.1169, 0.0833, 0.2562, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1759, 0.2309],\n",
      "        [0.1629, 0.0000, 0.1371, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.44200000166893005\n",
      "tensor([[ 0.0377,  0.0225,  0.0200,  0.0269,  0.0233],\n",
      "        [ 0.0387,  0.0793,  0.0615,  0.0604,  0.0616],\n",
      "        [-0.0047, -0.0050, -0.0115, -0.0015, -0.0040],\n",
      "        [-0.0025, -0.0054, -0.0046, -0.0025, -0.0038],\n",
      "        [-0.0052, -0.0070, -0.0058, -0.0054, -0.0045],\n",
      "        [ 0.0090,  0.0111,  0.0092,  0.0109,  0.0088],\n",
      "        [ 0.0047,  0.0066,  0.0051,  0.0073,  0.0062],\n",
      "        [ 0.0006, -0.0024, -0.0010, -0.0013, -0.0009],\n",
      "        [-0.0290, -0.0075, -0.0287, -0.0139, -0.0237],\n",
      "        [-0.0236, -0.0080, -0.0282, -0.0006, -0.0081]])\n",
      "tensor([ 0.0429,  0.1171, -0.0106, -0.0070, -0.0105,  0.0196,  0.0115, -0.0018,\n",
      "        -0.0410, -0.0283])\n",
      "tensor([[-0.0502, -0.0470, -0.0081,  0.0102,  0.0065,  0.0115,  0.0230,  0.0122,\n",
      "         -0.0768, -0.1012],\n",
      "        [-0.0582, -0.0302, -0.0025,  0.0128,  0.0067,  0.0115,  0.0216,  0.0096,\n",
      "         -0.0454, -0.0766],\n",
      "        [ 0.0088,  0.0024,  0.0011, -0.0017, -0.0010, -0.0013, -0.0016, -0.0008,\n",
      "         -0.0023,  0.0030],\n",
      "        [ 0.0618,  0.0234, -0.0007, -0.0138, -0.0077, -0.0135, -0.0292, -0.0121,\n",
      "          0.0588,  0.1158],\n",
      "        [ 0.0324,  0.0156,  0.0003, -0.0060, -0.0040, -0.0069, -0.0151, -0.0070,\n",
      "          0.0362,  0.0677],\n",
      "        [ 0.0109,  0.0115,  0.0017, -0.0025, -0.0015, -0.0033, -0.0075, -0.0034,\n",
      "          0.0237,  0.0340],\n",
      "        [ 0.0259,  0.0250,  0.0025, -0.0037, -0.0022, -0.0055, -0.0086, -0.0045,\n",
      "          0.0327,  0.0397],\n",
      "        [-0.0085, -0.0067, -0.0010,  0.0017,  0.0011,  0.0019,  0.0039,  0.0019,\n",
      "         -0.0114, -0.0169],\n",
      "        [ 0.0117,  0.0149,  0.0028, -0.0032, -0.0020, -0.0036, -0.0091, -0.0044,\n",
      "          0.0329,  0.0413],\n",
      "        [-0.0051, -0.0028, -0.0002,  0.0012,  0.0006,  0.0011,  0.0022,  0.0009,\n",
      "         -0.0048, -0.0081]])\n",
      "tensor([-0.2458, -0.2351,  0.0176,  0.2961,  0.1377,  0.0814,  0.0891, -0.0413,\n",
      "         0.0987, -0.0239])\n",
      "tensor([[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 1.1637e-01,  1.3983e-01,  9.1797e-02,  1.1349e-01,  1.2260e-02,\n",
      "          1.0266e-02,  9.1964e-02, -1.0624e-02,  7.1026e-02, -1.4382e-02],\n",
      "        [ 9.0686e-02,  1.2277e-01,  8.4020e-02,  9.4583e-02,  1.2332e-02,\n",
      "          9.8525e-03,  7.1444e-02, -8.7481e-03,  6.5071e-02, -1.2216e-02],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 7.7560e-02,  7.5569e-02,  5.8388e-02,  8.4299e-02,  5.8130e-03,\n",
      "          1.5624e-03,  8.6786e-02, -5.5060e-03,  3.0870e-02, -1.0580e-02],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 8.9482e-02,  1.1934e-01,  8.0998e-02,  9.2127e-02,  1.2411e-02,\n",
      "          8.3653e-03,  7.0879e-02, -8.6198e-03,  6.2769e-02, -1.1899e-02],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 3.8582e-02,  4.7337e-02,  3.6223e-02,  5.0127e-02,  5.6179e-03,\n",
      "          1.5867e-03,  4.4381e-02, -2.5404e-03,  1.9734e-02, -5.2773e-03],\n",
      "        [ 8.1950e-02,  5.2707e-02,  3.0884e-02,  6.5392e-02,  1.9031e-03,\n",
      "         -6.5221e-04,  7.7645e-02, -5.7908e-03,  1.6375e-02, -8.4566e-03],\n",
      "        [ 3.0286e-02,  2.3561e-02,  1.4224e-02,  2.8120e-02,  2.3804e-03,\n",
      "          1.0179e-03,  2.7486e-02, -2.0561e-03,  6.6299e-03, -3.1295e-03],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [-2.4299e-02, -1.2230e-02, -6.2974e-03, -2.2397e-02, -1.4592e-03,\n",
      "          1.6233e-03, -2.7907e-02,  1.1377e-03,  3.1805e-04,  2.0286e-03],\n",
      "        [ 6.3798e-02,  1.0283e-01,  7.2982e-02,  7.0282e-02,  1.0471e-02,\n",
      "          8.7573e-03,  4.8150e-02, -6.9716e-03,  6.0007e-02, -9.6681e-03],\n",
      "        [ 1.1738e-03,  4.4719e-04,  3.8843e-04,  1.0500e-03, -3.1734e-05,\n",
      "         -3.4269e-05,  1.4687e-03, -5.8484e-05, -8.0544e-07, -1.4034e-04],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [-3.6324e-02, -8.8178e-02, -6.4949e-02, -4.4848e-02, -8.4625e-03,\n",
      "         -1.0458e-02, -1.6749e-02,  5.6115e-03, -5.9764e-02,  7.2437e-03],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n",
      "          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00]])\n",
      "tensor([ 0.0000,  0.0000,  0.5950,  0.4950,  0.0000,  0.0000,  0.4000,  0.0000,\n",
      "         0.4850,  0.0000,  0.0000,  0.2095,  0.3450,  0.1300,  0.0000,  0.0000,\n",
      "         0.0000,  0.0000, -0.0920,  0.3850,  0.0050,  0.0000, -0.2780,  0.0000,\n",
      "         0.0000])\n",
      "epoch:  0 | ave_tr_loss:  2.86e-01 | te_loss:  3.17e-01 | acc.:  0.000000 | lr:  1.00e-04 | time:  104.453966     \n",
      "tensor([[0.1473, 0.0000, 0.0291, 0.1690, 0.0000],\n",
      "        [0.0000, 0.2114, 0.0000, 0.2236, 0.0000],\n",
      "        [0.0000, 0.2186, 0.1193, 0.2210, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2095, 0.2270],\n",
      "        [0.1851, 0.0000, 0.1997, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.44999998807907104\n",
      "tensor([[0.1482, 0.0000, 0.0136, 0.1584, 0.0000],\n",
      "        [0.0000, 0.1876, 0.0000, 0.2553, 0.0000],\n",
      "        [0.0000, 0.2365, 0.1439, 0.2124, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2104, 0.2390],\n",
      "        [0.1956, 0.0000, 0.2272, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.44519999623298645\n",
      "tensor([[0.1504, 0.0000, 0.0207, 0.1673, 0.0000],\n",
      "        [0.0000, 0.1972, 0.0000, 0.2260, 0.0000],\n",
      "        [0.0000, 0.2144, 0.1374, 0.2177, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2047, 0.2362],\n",
      "        [0.1822, 0.0000, 0.2034, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.43639999628067017\n",
      "tensor([[0.1447, 0.0000, 0.0164, 0.1703, 0.0000],\n",
      "        [0.0000, 0.2086, 0.0000, 0.2467, 0.0000],\n",
      "        [0.0000, 0.2378, 0.1254, 0.2185, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2070, 0.2350],\n",
      "        [0.1823, 0.0000, 0.2162, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.44920000433921814\n",
      "tensor([[0.1416, 0.0000, 0.0336, 0.1980, 0.0000],\n",
      "        [0.0000, 0.2243, 0.0000, 0.2200, 0.0000],\n",
      "        [0.0000, 0.1940, 0.0984, 0.2408, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1742, 0.2416],\n",
      "        [0.1628, 0.0000, 0.1917, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.4424000084400177\n",
      "tensor([[ 0.0369,  0.0219,  0.0198,  0.0277,  0.0210],\n",
      "        [ 0.0431,  0.0826,  0.0671,  0.0617,  0.0682],\n",
      "        [-0.0028, -0.0040, -0.0083, -0.0007, -0.0012],\n",
      "        [-0.0021, -0.0043, -0.0046, -0.0017, -0.0026],\n",
      "        [-0.0054, -0.0070, -0.0061, -0.0051, -0.0046],\n",
      "        [ 0.0094,  0.0115,  0.0095,  0.0105,  0.0090],\n",
      "        [ 0.0047,  0.0066,  0.0051,  0.0072,  0.0062],\n",
      "        [ 0.0007, -0.0022, -0.0010, -0.0009, -0.0012],\n",
      "        [-0.0240, -0.0079, -0.0265, -0.0122, -0.0175],\n",
      "        [-0.0246, -0.0094, -0.0298, -0.0028, -0.0050]])\n",
      "tensor([ 0.0430,  0.1251, -0.0073, -0.0057, -0.0106,  0.0198,  0.0113, -0.0016,\n",
      "        -0.0353, -0.0295])\n",
      "tensor([[-0.0533, -0.0495, -0.0082,  0.0104,  0.0065,  0.0123,  0.0238,  0.0124,\n",
      "         -0.0808, -0.1058],\n",
      "        [-0.0579, -0.0305, -0.0034,  0.0125,  0.0066,  0.0112,  0.0207,  0.0094,\n",
      "         -0.0439, -0.0733],\n",
      "        [ 0.0134,  0.0034,  0.0010, -0.0024, -0.0014, -0.0021, -0.0027, -0.0013,\n",
      "         -0.0013,  0.0071],\n",
      "        [ 0.0649,  0.0262, -0.0008, -0.0140, -0.0078, -0.0142, -0.0301, -0.0123,\n",
      "          0.0618,  0.1206],\n",
      "        [ 0.0325,  0.0167,  0.0005, -0.0058, -0.0037, -0.0073, -0.0151, -0.0068,\n",
      "          0.0382,  0.0692],\n",
      "        [ 0.0118,  0.0125,  0.0013, -0.0024, -0.0014, -0.0035, -0.0076, -0.0033,\n",
      "          0.0246,  0.0355],\n",
      "        [ 0.0260,  0.0263,  0.0033, -0.0043, -0.0025, -0.0055, -0.0089, -0.0048,\n",
      "          0.0355,  0.0393],\n",
      "        [-0.0088, -0.0071, -0.0010,  0.0017,  0.0010,  0.0021,  0.0040,  0.0019,\n",
      "         -0.0120, -0.0175],\n",
      "        [ 0.0110,  0.0141,  0.0019, -0.0029, -0.0017, -0.0035, -0.0087, -0.0041,\n",
      "          0.0325,  0.0406],\n",
      "        [-0.0051, -0.0027, -0.0002,  0.0012,  0.0006,  0.0011,  0.0021,  0.0009,\n",
      "         -0.0046, -0.0079]])\n",
      "tensor([-0.2522, -0.2236,  0.0286,  0.3016,  0.1389,  0.0801,  0.0940, -0.0419,\n",
      "         0.0930, -0.0228])\n",
      "tensor([[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.1203,  0.1476,  0.0956,  0.1167,  0.0121,  0.0098,  0.0930, -0.0109,\n",
      "          0.0746, -0.0147],\n",
      "        [ 0.0922,  0.1288,  0.0871,  0.0962,  0.0123,  0.0096,  0.0709, -0.0089,\n",
      "          0.0680, -0.0123],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0771,  0.0713,  0.0543,  0.0834,  0.0058,  0.0018,  0.0865, -0.0051,\n",
      "          0.0259, -0.0101],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0911,  0.1259,  0.0845,  0.0943,  0.0124,  0.0088,  0.0700, -0.0087,\n",
      "          0.0660, -0.0120],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0337,  0.0435,  0.0348,  0.0459,  0.0046,  0.0017,  0.0405, -0.0022,\n",
      "          0.0180, -0.0050],\n",
      "        [ 0.0776,  0.0470,  0.0265,  0.0608,  0.0012, -0.0008,  0.0736, -0.0053,\n",
      "          0.0129, -0.0077],\n",
      "        [ 0.0297,  0.0239,  0.0137,  0.0280,  0.0020,  0.0011,  0.0263, -0.0019,\n",
      "          0.0065, -0.0029],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [-0.0283, -0.0149, -0.0091, -0.0266, -0.0021,  0.0018, -0.0334,  0.0014,\n",
      "         -0.0003,  0.0027],\n",
      "        [ 0.0633,  0.1056,  0.0741,  0.0693,  0.0103,  0.0085,  0.0459, -0.0070,\n",
      "          0.0618, -0.0096],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [-0.0341, -0.0900, -0.0643, -0.0417, -0.0082, -0.0103, -0.0112,  0.0055,\n",
      "         -0.0612,  0.0067],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "          0.0000,  0.0000]])\n",
      "tensor([ 0.0000,  0.0000,  0.6140,  0.5050,  0.0000,  0.0000,  0.3850,  0.0000,\n",
      "         0.4950,  0.0000,  0.0000,  0.1910,  0.3190,  0.1250,  0.0000,  0.0000,\n",
      "         0.0000,  0.0000, -0.1135,  0.3850,  0.0000,  0.0000, -0.2665,  0.0000,\n",
      "         0.0000])\n",
      "epoch:  0 | ave_tr_loss:  3.22e-01 | te_loss:  3.17e-01 | acc.:  0.000000 | lr:  1.00e-04 | time:  130.623009     \n",
      "tensor([[0.1689, 0.0000, 0.0736, 0.2294, 0.0000],\n",
      "        [0.0000, 0.2959, 0.0000, 0.1024, 0.0000],\n",
      "        [0.0000, 0.0822, 0.0692, 0.2618, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1343, 0.2277],\n",
      "        [0.1025, 0.0000, 0.0746, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n",
      "0.44440001249313354\n",
      "tensor([[0.1444, 0.0000, 0.0527, 0.2121, 0.0000],\n",
      "        [0.0000, 0.2599, 0.0000, 0.1663, 0.0000],\n",
      "        [0.0000, 0.1630, 0.0728, 0.2504, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.1717, 0.2217],\n",
      "        [0.1416, 0.0000, 0.1473, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-18-36a83c23c81f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     20\u001b[0m         \u001b[0mweight_pred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md_batch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mweight_pred\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mweight_pred\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpath_batch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     23\u001b[0m         \u001b[0mpath_pred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdijkstra\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mweight_pred\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m         \u001b[0mloss2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath_pred\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpath_batch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1108\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1111\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1112\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Research/Implicit-Networks/SPO_with_DYS/source/fenchel_young.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_tensor, y_true, *args)\u001b[0m\n\u001b[1;32m     86\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     87\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput_tensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_true\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 88\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mPerturbedFunc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mapply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_tensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_true\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mperturbed\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_batched\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_maximize\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     89\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Research/Implicit-Networks/SPO_with_DYS/source/fenchel_young.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(ctx, input_tensor, y_true, perturbed, batched, maximize, *args)\u001b[0m\n\u001b[1;32m     31\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0mstaticmethod\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     32\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mctx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput_tensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_true\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mperturbed\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatched\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmaximize\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 33\u001b[0;31m         \u001b[0mdiff\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mperturbed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_tensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0my_true\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_tensor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     34\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mmaximize\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m             \u001b[0mdiff\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0mdiff\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Research/Implicit-Networks/SPO_with_DYS/source/perturbations.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(input_tensor, *args)\u001b[0m\n\u001b[1;32m    199\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moriginal_input_shape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    200\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 201\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mPerturbedFunc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mapply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_tensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    202\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    203\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Research/Implicit-Networks/SPO_with_DYS/source/perturbations.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(ctx, input_tensor, *args)\u001b[0m\n\u001b[1;32m    156\u001b[0m                 \u001b[0mperturbed_input\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mperturbed_input\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mflat_batch_dim_shape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    157\u001b[0m                 \u001b[0;31m# Calls user-defined function in a perturbation agnostic manner.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 158\u001b[0;31m                 \u001b[0mperturbed_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mperturbed_input\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    159\u001b[0m                 \u001b[0;31m# [NB, D1, ..., Dk] ->  [N, B, D1, ..., Dk].\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    160\u001b[0m                 \u001b[0mperturbed_input\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mperturbed_input\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mperturbed_input_shape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Research/Implicit-Networks/SPO_with_DYS/source/torch_Dijkstra.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, tensor)\u001b[0m\n\u001b[1;32m    106\u001b[0m                       axis=0)\n\u001b[1;32m    107\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m3\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 108\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    109\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_single\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    110\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Research/Implicit-Networks/SPO_with_DYS/source/torch_Dijkstra.py\u001b[0m in \u001b[0;36mrun_batch\u001b[0;34m(self, tensor)\u001b[0m\n\u001b[1;32m     96\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     97\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mrun_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 98\u001b[0;31m     return torch.stack([self.run_single(tensor[i])\n\u001b[0m\u001b[1;32m     99\u001b[0m                      for i in range(tensor.shape[0])],\n\u001b[1;32m    100\u001b[0m                     axis=0)\n",
      "\u001b[0;32m~/Research/Implicit-Networks/SPO_with_DYS/source/torch_Dijkstra.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     96\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     97\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mrun_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 98\u001b[0;31m     return torch.stack([self.run_single(tensor[i])\n\u001b[0m\u001b[1;32m     99\u001b[0m                      for i in range(tensor.shape[0])],\n\u001b[1;32m    100\u001b[0m                     axis=0)\n",
      "\u001b[0;32m~/Research/Implicit-Networks/SPO_with_DYS/source/torch_Dijkstra.py\u001b[0m in \u001b[0;36mrun_single\u001b[0;34m(self, costs, Gen_Data)\u001b[0m\n\u001b[1;32m     60\u001b[0m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcosts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     61\u001b[0m     \u001b[0;32mwhile\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mqueue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 62\u001b[0;31m       \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mheapq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mheappop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mqueue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     63\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvisits\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     64\u001b[0m         \u001b[0;32mcontinue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "## Train!\n",
    "import time\n",
    "from utils import Edge_to_Node, compute_perfect_path_acc, compute_perfect_path_acc_vertex\n",
    "max_epochs = 100\n",
    "max_time = 1e10\n",
    "test_loss_hist= []\n",
    "test_acc_hist = []\n",
    "train_time = []\n",
    "train_loss_ave = 0\n",
    "train_start_time = time.time()\n",
    "epoch=0\n",
    "epoch_time=0\n",
    "\n",
    "while epoch <= max_epochs and epoch_time <= max_time:\n",
    "    for d_batch, path_batch in train_loader:\n",
    "        d_batch = d_batch.to(device)\n",
    "        path_batch =path_batch.to(device)\n",
    "        net.train()\n",
    "        optimizer.zero_grad()\n",
    "        weight_pred = net(d_batch)\n",
    "        print(weight_pred[1,:,:])\n",
    "        loss = criterion(weight_pred, path_batch).mean()\n",
    "        path_pred = dijkstra(weight_pred)\n",
    "        loss2 = criterion2(path_pred, path_batch)\n",
    "        print(loss2.item())\n",
    "        train_loss_ave = 0.95*train_loss_ave + 0.05*loss2.item()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    # print('epoch:', epoch, ', av. training loss = ', train_loss_ave)\n",
    "    epoch_time = time.time() - train_start_time\n",
    "    train_time.append(epoch_time)\n",
    "\n",
    "    # Evaluate progress on test set. (note one batch is entire dataset)\n",
    "    net.eval()\n",
    "    for d_batch, path_batch in test_loader:\n",
    "        d_batch = d_batch.to(device)\n",
    "        path_batch =path_batch.to(device)\n",
    "        weight_pred = net(d_batch)\n",
    "        path_pred = dijkstra(weight_pred)\n",
    "        test_loss = criterion2(weight_pred, path_batch).mean().item()\n",
    "        test_loss_hist.append(test_loss)\n",
    "        for param in net.parameters():\n",
    "            print(param.grad)\n",
    "        # print('epoch: ', epoch, 'test loss is ', test_loss)\n",
    "        ## Evaluate accuracy\n",
    "        accuracy = compute_perfect_path_acc_vertex(path_pred, path_batch)\n",
    "        # regret = compute_regret(WW, d_batch, path_batch, path_pred,'V', Edge_list, grid_size, device)\n",
    "        # print('epoch: ', epoch, 'accuracy is ', accuracy)\n",
    "        test_acc_hist.append(accuracy)\n",
    "\n",
    "    # if test_loss < best_loss:\n",
    "    #     best_params = net.state_dict()\n",
    "\n",
    "    print('epoch: ', epoch, '| ave_tr_loss: ', \"{:5.2e}\".format(train_loss_ave), '| te_loss: ', \"{:5.2e}\".format(test_loss), '| acc.: ', \"{:<7f}\".format(accuracy), '| lr: ', \"{:5.2e}\".format(optimizer.param_groups[0]['lr']), '| time: ', \"{:<15f}\".format(epoch_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.1435, 0.0000, 0.0082, 0.1712, 0.0000],\n",
      "        [0.0000, 0.2040, 0.0000, 0.2673, 0.0000],\n",
      "        [0.0000, 0.2524, 0.1311, 0.2151, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.2081, 0.2429],\n",
      "        [0.1842, 0.0000, 0.2307, 0.0000, 0.0000]], grad_fn=<SliceBackward0>)\n"
     ]
    }
   ],
   "source": [
    "print(weight_pred[2,:,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 1., 0., 0., 0.],\n",
      "        [0., 1., 1., 0., 0.],\n",
      "        [0., 0., 1., 0., 0.],\n",
      "        [0., 0., 1., 0., 0.],\n",
      "        [0., 0., 1., 1., 1.]])\n"
     ]
    }
   ],
   "source": [
    "print(path_batch[2,:,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
